# -*- coding: utf-8 -*-
# coding=utf-8

import copy
from utils.log import logger
from transformers import AutoTokenizer, AutoModel

class LiteralRecognition:
    def __init__(self, model) -> None:
        self.model = model
        self.prompt = "你是一个字面量值提取器，输入一个句子以及一个字面量，需要输出该字面量对应的值和该值所涉及的运算符，注意不要输出解释性文本。输入：{}；输出：\n"

    def recognize(self, query, sel):
        """
        识别header对应的字面量
        :param query: 自然语言查询
        :param sel: 任务select封装的输出
        :return:一个literals列表，列表中的每个元素对应一个字典，元素的顺序也sel中出现的header一致
        """
        origin=[]
        literals = []
        for item in sel:
            header = item['name']
            query_header = query + ' ' + header
            # 模型预测
            try:
                prediction = self.model.predict(self.prompt.format(query_header))
                origin.append({header: prediction})
            except Exception as err:
                logger.error("literal prediction error happened![{}]: {}".format(type(err), str(err)))
                raise
            # 返回值去掉两边空格，然后以空格符为分隔符，获取一个列表
            new_dict = {}
            res_lst = prediction.strip().split()
            if len(res_lst) == 2:
                new_dict['literal'] = res_lst[0]
                new_dict['operator'] = res_lst[1]
            else:
                new_dict['literal'] = 'none'
                new_dict['operator'] = 'none'
            new_item = copy.deepcopy(item)
            new_item.update(new_dict)
            literals.append(new_item)
        return {'origin':origin,'processed':{"literals": literals}}

